#include "heap.h"

typedef struct HEAP_NODE {
  struct HEAP_NODE *lpPrev;
  struct HEAP_NODE *lpNext;
  DWORD dwSize;
} * PHEAP_NODE, *LPHEAP_NODE;

STATIC LPHEAP_NODE lpStart;
STATIC LPHEAP_NODE lpBefore;

VOID HeapInit(LPVOID lpvStart, DWORD dwSize) {
  lpStart = (LPHEAP_NODE)lpvStart;
  REG LPHEAP_NODE lpEnd =
      (LPHEAP_NODE)((DWORD_PTR)lpvStart + dwSize - sizeof(struct HEAP_NODE));
  lpEnd->dwSize = (DWORD_PTR)lpStart - (DWORD_PTR)lpEnd;
  lpEnd->lpPrev = (LPHEAP_NODE)lpvStart;
  lpEnd->lpNext = lpStart;
  lpStart->dwSize = sizeof(struct HEAP_NODE);
  lpStart->lpPrev = lpEnd;
  lpStart->lpNext = lpEnd;
  lpBefore = lpStart;
}

LPVOID HeapAlloc(DWORD dwSize) {
  REG LPHEAP_NODE lpNode;
  REG LPHEAP_NODE lpBlock;
  if (dwSize & 0x3)
    dwSize = (dwSize & 0xFFFFFFFC) + 0x4;
  dwSize += sizeof(struct HEAP_NODE);
  lpNode = lpBefore;
  do {
    if ((DWORD_PTR)lpNode->lpNext - (DWORD_PTR)lpNode >=
        dwSize + lpNode->dwSize) {
      lpBlock = (LPHEAP_NODE)((DWORD_PTR)lpNode + lpNode->dwSize);
      lpBlock->dwSize = dwSize;
      lpBlock->lpPrev = lpNode;
      lpBlock->lpNext = lpNode->lpNext;
      lpNode->lpNext->lpPrev = lpBlock;
      lpNode->lpNext = lpBlock;
      lpBefore = lpBlock;
      return (lpBlock + 1);
    }
    lpNode = lpNode->lpNext;
  } while (lpNode != lpBefore);
  return NULL;
}

VOID HeapFree(LPVOID lpvMemory) {
  ((LPHEAP_NODE)lpvMemory)->lpNext->lpPrev = ((LPHEAP_NODE)lpvMemory)->lpPrev;
  ((LPHEAP_NODE)lpvMemory)->lpPrev->lpNext = ((LPHEAP_NODE)lpvMemory)->lpNext;
}
